#!/usr/bin/env python

import os
import sys
import re
import random
from datetime import *
import time
import signal
import traceback
from sets import Set

def getChildIdList(sid):
    if sid not in idTree:
        return []
    else:
        childIdList = list(idTree[sid])
        graChildIdList = []
        for chid in childIdList:
            graChildIdList.extend(getChildIdList(chid))
        return childIdList+graChildIdList

timefmt='%Y-%m-%d %H:%M:%S'

#jobSign="pbssrv.ihep.ac.cn.SC"
selfpid=os.getpid()

log_file_path=sys.path[0]
log_file_name=sys.argv[0]
log_file_name=os.path.basename(log_file_name)
log_file_name=os.path.splitext(log_file_name)[0]
hostinfo=os.uname()
hostname=hostinfo[1]
random.seed(hostname)
log_file_name="/tmp/"+log_file_name+'_'+hostname+'_'+\
              time.strftime('%Y_%m_%d')+".log"

log_file=open(log_file_name,"a")
log_file.write(time.strftime(timefmt)+" START cleaning and OPEN log file\n")
single_file_name="/usr/local/zombie.log"
single_file=open(single_file_name,"w")
log_file.write(time.strftime(timefmt)+" Get PBSNODES status\n")

subproc=os.popen("ps -ejl").readlines()
#time.sleep(90)
log_file.write(time.strftime(timefmt)+" Get PROCESS status\n")
#subproc.read()
#allproc=subproc.read()

################### change this part  ###############################
##subproc1=os.popen("cat /home/cc/suzz/gpuqstat/gpuqstat | grep '"\
##		 +re.split('\.',hostname)[0]+"\S*$'")
##jsidstr=subproc1.readlines()

##subproc1.close()
##jsidstr=[re.sub('\[\d+\]', '', i) for i in jsidstr]
##log_file.write(time.strftime(timefmt)+" Get QSTAT status\n")

##jsidstr=[re.match("^([0-9]+)\.\S+\s+\S+\s+\S+\s+\S+\s+(\S+)",i).groups() \
##		for i in jsidstr]
##jsidstr=[[int(j) for j in i] for i in jsidstr]
##jsiddict=dict(jsidstr)

##jobsIdAll = open("/home/cc/suzz/gpuqstat/gpuqstat").readlines()

##jobsIds = []
##for jobsId in jobsIdAll:
##	jobsId = jobsId.split(".")[0]
##	jobsIds.append(jobsId)

######################## change part  ##################################
jobsIdAll = open("/home/bes/suzz/gpuqstat/pbssrv02stat").readlines()

jsiddict={}
jobsIds = []
host = os.uname()[1].split(".")[0]
if (len(jobsIdAll)<3):
    log_file.write("failed to get  pbssrv02qstat ")
    sys.exit(1) 
##change by Kan Bowen
#f = os.sys.argv[0]
f = '/home/bes/suzz/gpuqstat/pbssrv02stat'
mtime = time.ctime(os.path.getmtime(f))
ctime = time.ctime(os.path.getctime(f))
print "Last modified : %s, last created time: %s" % (mtime, ctime)
print os.path.getctime(f)
print "Now: %s " %time.time()
creattime = os.path.getctime(f)
nowtime = time.time()
print int((creattime+360))-int(nowtime)
if (int((creattime+180))<=int(nowtime)):
    log_file.write("failed to get  pbssrv02stat ")
    sys.exit(1)


##changed by Kan Bowen

for jobsId in jobsIdAll:
	if host in jobsId:
		parts=jobsId.split()
		sessionId=parts[4]
		hostjobid=parts[0].split(".")[0]
		jsiddict[hostjobid]= int(sessionId)
	jobsId = jobsId.split(".")[0]
	jobsIds.append(jobsId)
log_file.write("jsiddict::::::\n")
log_file.write(str(jsiddict))
lastJobsId = jobsIds[-1]
log_file.write("lastJobId="+str(lastJobsId)+"\n")
#jsiddict2 = {}
for jobid in jsiddict.keys():
	if int(jobid) <= int(lastJobsId):
		if str(jobid) not in jobsIds:
			jsiddict.pop(jobid)
			#print "delete jobid:"
			#print jobid

	
#	if int(jobid) >= int(lastJobsId):
#		subproc2=os.popen("qstat -an1 -t | grep '"\
#                 +re.split('\.',hostname)[0]+"\S*$'")
#		jsidstr2=subproc2.readlines()
#
#		subproc2.close()
#		jsidstr2=[re.sub('\[\d+\]', '', i) for i in jsidstr2]
#		log_file.write(time.strftime(timefmt)+" Get QSTAT status\n")
#
#		jsidstr2=[re.match("^([0-9]+)\.\S+\s+\S+\s+\S+\s+\S+\s+(\S+)",i).groups() \
#                	for i in jsidstr2]
#		jsidstr2=[[int(j) for j in i] for i in jsidstr2]
#		jsiddict2=dict(jsidstr2)
#print jsiddict2
#jsiddict.update(jsiddict2)
white_sid=jsiddict.values()
white_jobids=jsiddict.keys()

log_file.write("The first:\n")
log_file.write(str(white_sid))
#print white_sid


psidstr=[i.strip() for i in subproc]
#print psidstr
#print white_cmd
#subproc.close()
#while '' in psidstr:
#	psidstr.remove('')
psidstr=psidstr[1:]
####psidstr=[re.split('\s+',i)[2:7]+\
####         [re.search('(\S+\s+){15}(.+)',i).group(2)] for i in psidstr]
psidstr=[re.split('\s+',i)[2:7]+\
         [re.search('(\S+\s+){15}(.+)',i).group(2)] for i in psidstr]
white_uid=[4294967294];
white_name=['nagios','ganglia']
for i in white_name:
        cmd = "cat /etc/passwd | grep %s | awk -F':' '{print $3}'" %i
        uid = os.popen(cmd)
        white_uid.extend(uid.readline().split())
	uid.close()

white_cmd="(fs[ -_]?)|(daemon ?)|(portmap ?)|(nrpe ?)"
killlist=[]


mpd=os.popen("ps -ef | grep mpd | grep -v grep | awk -F' ' '{print $2}'")
mpdProcess=[int(s.strip()) for s in mpd.readlines()]
mpd.close()

pluto=os.popen("ps -ef | grep pluto | grep -v grep | awk -F' ' '{print $2}'")
plutoProcess=[int(s.strip()) for s in pluto.readlines()]
pluto.close()

#pbsmom=os.popen("ps -ef | grep pbs_mom | grep -v grep | awk -F' ' '{print $2}'")
#momProcess=[int(s.strip()) for s in pbsmom.readlines()]
#pbsmom.close()

#mompriv=os.popen("ps -ef | grep mom_priv | grep -v grep | awk -F' ' '{print $2}'")
#momprivProcess=[int(s.strip()) for s in mompriv.readlines()]
#mompriv.close()

#pbsJobs=os.popen("ps -ef|grep -v grep |grep "+jobSign).readlines()
whiteJobProcess=[]
#for eachJob in pbsJobs:
#    finalstr=eachJob.split()
#    #print finalstr
#    if (len(finalstr)>8):
#       lastCom=finalstr[8].split("/")[-1]
#       if (jobSign in lastCom):
#            checkJobId=int((lastCom.split(".")[0]))
###	    lastJobsId="10602150"
#            print "checkJobId=%d, lastJobsId=%s\n" %(checkJobId,lastJobsId)
#            if (checkJobId > int(lastJobsId)):
#		whiteJobProcess.append(int(finalstr[1]))

idTree = {}
for i in psidstr:
    chid = int(i[1])
    psid = int(i[2])
    if psid in idTree:
        idTree[psid].add(chid)
    else:
        idTree[psid] = Set([chid])

##for each in idTree:
##    print "idTree["+str(each)+"]="+str(idTree[each])
    

whiteIdList = []

#white_sid=white_sid+mpdProcess
white_sid=list(Set(white_sid).union(mpdProcess))
white_sid=list(Set(white_sid).union(plutoProcess))
#white_sid=list(Set(white_sid).union(whiteJobProcess))
#white_sid=list(Set(white_sid).union(momProcess))
#white_sid=list(Set(white_sid).union(momprivProcess))

for sid in white_sid:
    whiteIdList.append(sid)
    whiteIdList.extend(getChildIdList(sid))

print "whiteIdList============="
print whiteIdList
print "whiteIdList--------------"

log_file.write("The whiteIdList:\n")
log_file.write(str(whiteIdList))


#print white_sid
#print white_uid
#print white_cmd
for i in psidstr:
    if i[0]!='0' and int(i[4]) not in white_sid:


        log_file.write(time.strftime(timefmt)+" NON-ROOT Process "+\
	                       "ID:"+i[1]+" ParentID:"+i[2]+" CMD:"+i[5]+\
                       " has ILLEGAL SESSION ID:"+i[4]+"\n")
        if int(i[0])<10000 or str(i[0]) in white_uid:
            log_file.write("            UID="+i[0]+" which is in the whitelist, check it manually.\n")
        elif re.search(white_cmd,i[5]):
            log_file.write("            Command is in the whitelist, check it manually.\n")
        else:


	    my_file_name=sys.argv[0]
	    #print my_file_name
	    #local=os.popen("ps aux | grep zombie_gpu.py | grep -v grep ")
            local=os.popen("ps aux | grep %s  | grep -v grep " %my_file_name) 
            localProcess=local.readline()
            local.close()
            localProcess=re.split('\s+',localProcess.strip())
            if int(localProcess[1])!=int(i[1]) and int(i[1]) not in whiteIdList and int(i[1]) not in mpdProcess:
		
		pbsmom=os.popen("ps -ef | grep pbs_mom | grep -v grep | awk -F' ' '{print $2}'").readline()


                if ((int(i[2])<>int(pbsmom) and str(i[5])<>'tcsh') or (int(i[2])<>int(pbsmom) and str(i[5]) <> 'bash')):

            
                        ####blankstr=os.popen("ps -eo pid,tty,user,comm,stime,etime | grep " + str(i[1]))
                        #blankstr=os.popen("ps -ef | grep " + str(i[1]))
                        #blankstr1=blankstr.readline()
                        #blankstr.close()


                        #finalstr=re.split('\s+',blankstr1.strip())


                        ####if "-" in finalstr[5]:
                        ####if int(finalstr[5].split('-')[1].split(':')[1])>1:


			### Edited by Kan Bowen
			#os.popen("qstat -an1 | grep ")
			newProcess=os.popen("qstat -an1 | grep '"+re.split('\.',hostname)[0]+"'")

			newProcessContext=newProcess.readlines()
			newProcess.close()
			newProcessContext=[re.sub('\[\d+\]', '', m) for m in newProcessContext]
			print newProcessContext

			newProcessContext=[re.match("^([0-9]+)\.\S+\s+\S+\s+\S+\s+\S+\s+(\S+)",m).groups() for m in newProcessContext]
        	       	newProcessContext=[[int(n) for n in m] for m in newProcessContext]
	               	newProcessContext=dict(newProcessContext)
			newSessionId=newProcessContext.values()
			newSessionJobs=newProcessContext.keys()
			new_whiteIdList=[]
			for sid in newSessionId:
    				new_whiteIdList.append(sid)
    				new_whiteIdList.extend(getChildIdList(sid))
			if int(i[1]) not in new_whiteIdList:
			###Edited by Kan Bowen
				try:
                                		os.kill(int(i[1]),9)
                                		print "kill: %d"%(int(i[1]))
                                		#print signal.SIGTERM
						log_file.write("All process list:\n")
						for j in psidstr:
							log_file.write(str(j)+"\n")
						log_file.write("start to write job in this node\n")
						for j in jsiddict:
							log_file.write(str(j)+"\n")

						for j in jobsIdAll:
							log_file.write(str(j)+"\n")

                      		except:
                                		traceback.print_exc(file=log_file)

                        			log_file.write("            signal SIGTERM sent.\n")

                       		killlist.append(i[1])
		
                if (int(i[2])==int(pbsmom) and str(i[5])=='tcsh'):
                        timefmt2='%Y-%m-%d'
                        nodeName=os.popen("hostname").readline().strip()
                        log_file_suspicious_path="/home/cc/kanbw/zombie/%s+%s" %(time.strftime(timefmt2),nodeName)
                        log_file_suspicious=open(log_file_suspicious_path,"a")

                        log_file_suspicious.write("The suspicious process:       \n")
                        log_file_suspicious.write("Name:%s, PID:%s, PPID:%s \n" %(i[5],i[1],i[2]))
                        log_file_suspicious.close()
                if (int(i[2])==int(pbsmom) and str(i[5])=='bash'):
                        timefmt2='%Y-%m-%d'
                        nodeName=os.popen("hostname").readline().strip()
                        log_file_suspicious_path="/home/cc/kanbw/zombie/%s+%s" %(time.strftime(timefmt2),nodeName)
                        log_file_suspicious=open(log_file_suspicious_path,"a")

                        log_file_suspicious.write("The suspicious process:       \n")
                        log_file_suspicious.write("Name:%s, PID:%s, PPID:%s \n" %(i[5],i[1],i[2]))
                        log_file_suspicious.close()


#time.sleep(20)
subproc=os.popen("ps -efj")
log_file.write(time.strftime(timefmt)+" Get PROCESS status\n")
psidstr=subproc.readlines()
subproc.close()
psidstr=psidstr[1:]
psidstr=[re.split('\s+',i)[1] for i in psidstr]
if len(killlist)== 0:
    single_file.write("OK;\n")
tag=0
zombieList=[]
for i in killlist:
    if i in psidstr:
        log_file.write(time.strftime(timefmt)+" Process ID:"+i+\
                       " is still ALIVE, or this PID is just reused.\n")
    	log_file.write(" Ths killlist: " + str(killlist)+ ".\n")
    	tag=8888
    	zombieList.append(i)
    else:
        tag=9999
        log_file.write(time.strftime(timefmt)+" Process ID:"+i+\
                       " successfully vaporized.\n")
                      # Talking about this process any more
                      # should be considered CRIMETHINK.
    log_file.write(" Ths killlist: " + str(killlist)+ ".\n")
zombieStr=str(zombieList)
zombieStr=zombieStr.replace('\'','')
if tag==8888:
    single_file.write("ERROR; " +str(len(zombieList))+ " jobs:" + zombieStr  + " zombie process which need check manually.\n")
if tag==9999:
    single_file.write("OK;\n")
log_file.write(time.strftime(timefmt)+" FINISHED\n\n")
log_file.close()


